package org.example.ai_chat.repository;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.*;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 重写JdbcChatMemoryRepository和MysqlChatMemoryRepository中的方法，使其能够支持自定义存储，默认存储表chat_ai_memory
 */
public class CustomTableMysqlChatMemoryRepository implements ChatMemoryRepository {

    private final JdbcTemplate jdbcTemplate;
    private final String tableName;

    public CustomTableMysqlChatMemoryRepository(JdbcTemplate jdbcTemplate, String tableName) {
        Assert.notNull(jdbcTemplate, "jdbcTemplate cannot be null");
        Assert.hasText(tableName, "tableName cannot be null or empty");
        this.jdbcTemplate = jdbcTemplate;
        this.tableName = tableName;
        this.checkAndCreateTable();
    }

    /**
     * 检查数据库中是否存在这个表，不存在就创建
     */
    private void checkAndCreateTable() {
        String checkSQL = String.format("select table_name FROM information_schema.tables where table_schema = database() and table_name = '%s'", this.tableName);
        String constraintName = this.tableName + "_chk_message_type";
        String createSQL = String.format("create table %s (id BIGINT auto_increment primary key, conversation_id varchar(256) not null, content longtext not null, type varchar(100) not null, timestamp timestamp not null, constraint %s check (type IN ('USER', 'ASSISTANT', 'SYSTEM', 'TOOL')))", this.tableName, constraintName);
        if (!(Boolean) jdbcTemplate.query(checkSQL, ResultSet::next)) {
            jdbcTemplate.execute(createSQL);
        }
    }


    @Override
    public List<String> findConversationIds() {
        String sql = String.format("select distinct conversation_id from %s", this.tableName);
        return jdbcTemplate.query(sql, (rs, rowNum) -> rs.getString(1));
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        String sql = String.format("select content, type from %s where conversation_id = ? order by timestamp", this.tableName);
        return jdbcTemplate.query(sql, new MessageRowMapper(), conversationId);
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        Assert.notNull(messages, "messages cannot be null");
        Assert.noNullElements(messages, "messages cannot contain null elements");

        this.deleteByConversationId(conversationId);
        String sql = String.format("insert into %s (conversation_id, content, type, timestamp) values (?, ?, ?, ?)", this.tableName);

        this.jdbcTemplate.batchUpdate(sql, new BatchPreparedStatementSetter() {
            private final AtomicLong instantSeq = new AtomicLong(System.currentTimeMillis());

            @Override
            public void setValues(PreparedStatement ps, int i) throws SQLException {
                Message message = messages.get(i);
                ps.setString(1, conversationId);
                ps.setString(2, message.getText());
                ps.setString(3, message.getMessageType().name());
                ps.setTimestamp(4, new Timestamp(instantSeq.getAndIncrement()));
            }

            @Override
            public int getBatchSize() {
                return messages.size();
            }
        });
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        String sql = String.format("delete from %s where conversation_id = ?", this.tableName);
        jdbcTemplate.update(sql, conversationId);
    }

    // Builder模式
    public static CustomBuilder builder() {
        return new CustomBuilder();
    }

    public static class CustomBuilder {
        private JdbcTemplate jdbcTemplate;
        private String tableName;

        public CustomBuilder jdbcTemplate(JdbcTemplate jdbcTemplate) {
            this.jdbcTemplate = jdbcTemplate;
            return this;
        }

        public CustomBuilder tableName(String tableName) {
            this.tableName = tableName;
            return this;
        }

        public CustomTableMysqlChatMemoryRepository build() {
            return new CustomTableMysqlChatMemoryRepository(jdbcTemplate, tableName);
        }
    }

    private static class MessageRowMapper implements RowMapper<Message> {
        private MessageRowMapper() {
        }

        @Nullable
        public Message mapRow(ResultSet rs, int i) throws SQLException {
            String content = rs.getString(1);
            MessageType type = MessageType.valueOf(rs.getString(2));
            Object var10000;
            switch (type) {
                case USER -> var10000 = new UserMessage(content);
                case ASSISTANT -> var10000 = new AssistantMessage(content);
                case SYSTEM -> var10000 = new SystemMessage(content);
                case TOOL -> var10000 = new ToolResponseMessage(List.of());
                default -> throw new IncompatibleClassChangeError();
            }

            return (Message)var10000;
        }
    }
}
